#!/usr/bin/python
from rap_core import *


zeff=10
core=-10.
ip=0.5


clevs=rap_states_ls()
clevs.set_nuclear_charge(zeff)
clevs.set_core_energy(core)
clevs.set_ionization_threshold(ip)

clevs.add_level('1s2','1S',0.0)
clevs.add_level('1s 2s','3S',0.09)
clevs.add_level('1s 2s','1S',0.11)
clevs.add_level('1s 2p','3P',0.14)
clevs.add_level('1s 2p','1P',0.15)

if clevs.nuclear_charge() != zeff:
  print 'problem with nuclear_charge()'
if clevs.core_energy() != core:
  print 'problem with core_energy()'
if clevs.ionization_threshold() != ip:
  print 'problem with ionization_threshold()'
if len(clevs) != 5:
  print 'wrong number of levels'

if str(clevs[1]) != '1s2 1S':
  print 'A - incorrect level: 1s2 1S'
if str(clevs[5]) != '1s1 2p1 1P':
  print 'A - incorrect level: 1s 2p 1P'


toadd=[]
toadd.append((0.00,2,'1s2','1S'))
toadd.append((0.09,3,'1s 2s','3S'))
toadd.append((0.11,4,'1s 2s','1S'))
toadd.append((0.14,5,'1s 2p','3P'))
toadd.append((0.15,6,'1s 2p','1P'))

clevs2=clevs.copy()
clevs2.clear_levels()
clevs2.add_levels(toadd)

if clevs2.nuclear_charge() != zeff:
  print 'problem with nuclear_charge()'
if clevs2.core_energy() != core:
  print 'problem with core_energy()'
if clevs2.ionization_threshold() != ip:
  print 'problem with ionization_threshold()'
if len(clevs2) != 5:
  print 'wrong number of levels'

if str(clevs2[2]) != '1s2 1S':
  print 'B - incorrect level: 1s2 1S'
if str(clevs2[6]) != '1s1 2p1 1P':
  print 'B - incorrect level: 1s 2p 1P'

clevs.set_lmap(clevs2)
if clevs.label(3,imap=True,spacer='.') != '1s1.2s1.3S':
  print 'problem with mapping or label'

clevs2.reindex()
if str(clevs2[1]) != '1s2 1S':
  print 'C - incorrect level: 1s2 1S'
if str(clevs2[5]) != '1s1 2p1 1P':
  print 'C - incorrect level: 1s 2p 1P'


levs=rap_states_ls('olg.v21')
idxc=levs.cfilter('2s 2p3')
idxt=levs.tfilter('3P',subset=idxc)
if idxc != [3,5,6,7,8,9]:
  print 'problem with cfilter()'
if idxt != [6]:
  print 'problem with tfilter()'


levs=rap_states_ls('ls#fe24.dat')
if len(levs) != 19:
  print 'problem opening adf04 file'
